#include <string.h>
#include <unistd.h>
#include <string>
#include <atomic>
#include <fstream>
#include <iostream>
#include <csignal>
#include <thread>  
#include <vector>
#include "rknn_api.h"
#include "rkllm.h"
#include <opencv2/opencv.hpp>

#define PROMPT_TEXT_PREFIX "<用户><image>"
#define PROMPT_TEXT_POSTFIX "\n<AI>"
#define LLM_HIDDEN_SIZE 2304
#define IMAGE_SIZE 384


using namespace std;

rknn_context ctx;
LLMHandle llmHandle = nullptr;


cv::Mat latest_frame;
bool RUN_FLAG = true;
std::mutex frame_mutex;

void exit_handler(int signal) {
    if (llmHandle != nullptr) {
        cout << "程序即将退出" << endl;
        LLMHandle _tmp = llmHandle;
        llmHandle = nullptr;
        rkllm_destroy(_tmp);
    }
    rknn_destroy(ctx);
    exit(signal);
}

void callback(RKLLMResult *result, void *userdata, LLMCallState state) {
    if (state == RKLLM_RUN_FINISH) {
        printf("\n");
    } else if (state == RKLLM_RUN_ERROR) {
        printf("\\run error\n");
    } else if (state == RKLLM_RUN_GET_LAST_HIDDEN_LAYER) {
        if (result->last_hidden_layer.embd_size != 0 && result->last_hidden_layer.num_tokens != 0) {
            int data_size = result->last_hidden_layer.embd_size * result->last_hidden_layer.num_tokens * sizeof(float);
            printf("\ndata_size:%d", data_size);
            std::ofstream outFile("last_hidden_layer.bin", std::ios::binary);
            if (outFile.is_open()) {
                outFile.write(reinterpret_cast<const char*>(result->last_hidden_layer.hidden_states), data_size);
                outFile.close();
                std::cout << "Data saved to output.bin successfully!" << std::endl;
            } else {
                std::cerr << "Failed to open the file for writing!" << std::endl;
            }
        }
    } else if (state == RKLLM_RUN_NORMAL) {
        printf("%s", result->text);
    }
}

// 读取视频帧的线程函数
void readFrames(cv::VideoCapture &cap) {
    cv::Mat frame;
    while (RUN_FLAG) {
        cap >> frame;
        if (frame.empty()) {
            std::cerr << "摄像头读取失败!" << std::endl;
            break;
        }
        // 获取图像的宽度和高度
        int h = frame.rows;
        int w = frame.cols;
        // 计算缩放比例，使最大边缩放到384
        float scale = float(IMAGE_SIZE) / std::max(h, w);
        int new_w = static_cast<int>(w * scale);
        int new_h = static_cast<int>(h * scale);

        cv::resize(frame, frame, cv::Size(new_w, new_h));
        // 创建一个 384x384 的黑色背景
        cv::Mat result_frame = cv::Mat::zeros(IMAGE_SIZE, IMAGE_SIZE, CV_8UC3);
        // 计算将缩放图像粘贴到背景的偏移量（居中）
        int x_offset = (384 - new_w) / 2;
        int y_offset = (384 - new_h) / 2;
        // 将缩放后的图像粘贴到 result_frame 中
        frame.copyTo(result_frame(cv::Rect(x_offset, y_offset, new_w, new_h)));


        latest_frame = result_frame;
    }
}

void displayFrames() {
    while (RUN_FLAG) {
        cv::Mat frame_copy;
        if (latest_frame.empty()) {
            continue;
        }
        frame_copy = latest_frame.clone();
        cv::imshow("img", frame_copy);
        cv::waitKey(1);
    }
}


int main() {
    signal(SIGINT, exit_handler);
    
    // =====================================rkllm model=============================
    // Initialize rkllm model
    RKLLMParam param = rkllm_createDefaultParam();
    param.model_path = "../../model.rkllm";
    param.top_k = 1;
    param.max_new_tokens = 256;
    param.max_context_len = 512;
    param.skip_special_token = true;
    param.img_start = "<image>";
    param.img_end = "</image>\n";
    param.img_content = "<unk>";
    int ret = rkllm_init(&llmHandle, &param, callback);
    if (ret != 0) {
        cout << "rkllm 模型加载失败\n";
        exit_handler(-1);
    }
    
    RKLLMInferParam rkllm_infer_params;
    memset(&rkllm_infer_params, 0, sizeof(RKLLMInferParam));
    rkllm_infer_params.mode = RKLLM_INFER_GENERATE;
    
    // =====================================rknn model=============================
    // Initialize rknn model
    string vmPath = "../../vision-384.rknn";
    int init_ret = rknn_init(&ctx, const_cast<char*>(vmPath.c_str()), 0, 0, nullptr);
    if (init_ret < 0) {
        cout << "RKNN 初始化失败。" << init_ret << endl;
        return -1;
    }

    rknn_input input;
    memset(&input, 0, sizeof(input));
    input.index = 0;  // 输入索引
    input.type = RKNN_TENSOR_UINT8;
    input.buf = nullptr;  // 输入数据
    input.size = 3 * IMAGE_SIZE * IMAGE_SIZE;  // 以字节为单位设置输入大小
    input.pass_through = false;  // 通常设置为false，表示不使用通过模式
    input.fmt = RKNN_TENSOR_NHWC;
    
    rknn_output output;
    memset(&output, 0, sizeof(output));
    output.index = 0;  // 输入索引
    output.buf = nullptr;  // 输入数据
    output.size = 64*LLM_HIDDEN_SIZE*sizeof(float);  // 以字节为单位设置输入大小
    output.is_prealloc = false;  // 通常设置为false，表示不使用通过模式
    output.want_float = true;


    // 摄像头
    string camUrl = "http://192.168.1.131:4747/video";
    cv::VideoCapture cap(camUrl);
    cap.set(cv::CAP_PROP_FOURCC, cv::VideoWriter::fourcc('M', 'J', 'P', 'G'));
    cap.set(cv::CAP_PROP_FRAME_WIDTH, 640);
    cap.set(cv::CAP_PROP_FRAME_HEIGHT, 480);
    cap.set(cv::CAP_PROP_FPS, 30);




    if (!cap.isOpened()) {
        std::cerr << "无法连接到摄像头" << std::endl;
        return -1;
    }
    std::thread clean_cache_capture_thread(readFrames, ref(cap));   // 清楚缓存的线程
    std::thread display_capture_thread(displayFrames);


    cout << "**********************现在，让我们开始对话吧。********************" << endl;
    while (true){
        string input_str;
        cout << "用户 (输入exit退出): ";
        getline(cin, input_str);
        if (input_str == "exit"){
            RUN_FLAG = false;
            break;
        }

        RKLLMInput rkllm_input;
        input_str = PROMPT_TEXT_PREFIX  + input_str + PROMPT_TEXT_POSTFIX;
        rkllm_input.input_type = RKLLM_INPUT_MULTIMODAL;
        rkllm_input.multimodal_input.prompt = (char*)input_str.c_str();
        rkllm_input.multimodal_input.n_image_tokens = 64;
        cv::Mat image = latest_frame.clone();
        cout<<"AI: ";
        cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
        input.buf = image.data;
        cv::imwrite("res.jpg", image);
        
        rknn_inputs_set(ctx, 1, &input);
        rknn_run(ctx, nullptr);
        rknn_outputs_get(ctx, 1, &output, nullptr);
        rkllm_input.multimodal_input.image_embed = (float *) output.buf;
        
        
        rkllm_run(llmHandle, &rkllm_input, &rkllm_infer_params, nullptr);

    }

    clean_cache_capture_thread.join();
    display_capture_thread.join();
    rkllm_destroy(llmHandle);
    return 0;
}
